package mr

import (
	"log"
	"net"
	"net/http"
	"net/rpc"
	"os"
	"sync"
	"time"
)

type Coordinator struct {
	// Your definitions here.
	mapList         []TaskAssignment
	reduceList      []TaskAssignment
	restM, restR, R int
	doneList        map[string]bool
	lock            sync.RWMutex
	cond            *sync.Cond
}

func (c *Coordinator) init(M, R int, files []string) {
	c.mapList = make([]TaskAssignment, M)
	for idx, filename := range files {
		c.mapList[idx] = TaskAssignment{Id: idx, Type: Map, DataPath: []string{filename}, NReduce: R}
	}

	c.reduceList = make([]TaskAssignment, R)
	for idx := range R {
		c.reduceList[idx] = TaskAssignment{Id: idx, Type: Reduce, DataPath: []string{}, NReduce: R}
	}

	c.restM, c.restR, c.R = M, R, R
	c.doneList = make(map[string]bool)
	c.cond = sync.NewCond(&c.lock)
}

// Your code here -- RPC handlers for the worker to call.

// an example RPC handler.
//
// the RPC argument and reply types are defined in rpc.go.
func (c *Coordinator) Example(args *ExampleArgs, reply *ExampleReply) error {
	reply.Y = args.X + 1
	return nil
}

// start a thread that listens for RPCs from worker.go
func (c *Coordinator) server() {
	rpc.Register(c)
	rpc.HandleHTTP()
	//l, e := net.Listen("tcp", ":1234")
	sockname := coordinatorSock()
	os.Remove(sockname)
	l, e := net.Listen("unix", sockname)
	if e != nil {
		log.Fatal("listen error:", e)
	}
	go http.Serve(l, nil)
}

// main/mrcoordinator.go calls Done() periodically to find out
// if the entire job has finished.
func (c *Coordinator) Done() bool {
	// ret := false

	// Your code here.

	// return ret
	c.lock.RLock()
	defer c.lock.RUnlock()
	return c.restR == 0
}

func (c *Coordinator) AssignTask(notUsed TaskResult, task *TaskAssignment) error {

	defer func() {
		go c.checkTimeout(*task)
	}()

	// All map tasks must be done before doing reduce tasks.
	if c.assignSpecificTask(&(c.restM), &(c.mapList), task) {
		return nil
	}

	// All map tasks are done. Now workers can do reduce tasks.
	// All reduce tasks must be done before exit.
	if c.assignSpecificTask(&(c.restR), &(c.reduceList), task) {
		return nil
	}

	// All tasks are done. The only task is to exit.
	*task = TaskAssignment{Id: 0, Type: Exit, DataPath: []string{""}, NReduce: 0}
	return nil
}

// If nRest > 0, assign task in taskList. Return true if assignment is successful.
func (c *Coordinator) assignSpecificTask(nRest *int, taskList *[]TaskAssignment, task *TaskAssignment) bool {
	c.lock.Lock()
	defer c.lock.Unlock()

	for *nRest > 0 {
		if length := len(*taskList); length > 0 {
			*task = (*taskList)[length-1]
			*taskList = (*taskList)[:length-1]
			return true
		} else {
			c.cond.Wait()
		}
	}
	return false
}

func (c *Coordinator) checkTimeout(task TaskAssignment) {
	if task.Type == Exit {
		return
	}

	time.Sleep(TIMEOUT)

	uid := task.Uid()

	c.lock.RLock()
	isDone := c.doneList[uid]
	c.lock.RUnlock()

	if isDone {
		return
	} else {
		c.ReturnTaskResult(task.ToTaskResult(), &TaskAssignment{})
	}

}

func (c *Coordinator) addTaskAssignment(task TaskAssignment) {
	var targetList *[]TaskAssignment
	if task.Type == Map {
		targetList = &(c.mapList)
	} else {
		targetList = &(c.reduceList)
	}

	uid := task.Uid()

	c.lock.Lock()
	defer c.lock.Unlock()
	*targetList = append(*targetList, task)
	c.doneList[uid] = false
}

func (c *Coordinator) ReturnTaskResult(taskResult TaskResult, notUsed *TaskAssignment) error {
	uid := taskResult.Uid()
	// if the task has been done, ignore it
	c.lock.RLock()
	if c.doneList[uid] {
		c.lock.RUnlock()
		return nil
	}
	c.lock.RUnlock()
	defer c.cond.Broadcast()

	// The worker failed to finish the task.
	if len(taskResult.ResultPath) == 0 {
		task := TaskAssignment{Id: taskResult.Id, Type: taskResult.Type, DataPath: taskResult.DataPath, NReduce: c.R}
		c.addTaskAssignment(task)
		return nil
	}

	// The worker success to finish the task.
	c.lock.Lock()
	defer c.lock.Unlock()

	if taskResult.Type == Map {
		for idx, filepath := range taskResult.ResultPath {
			if filepath != "" {
				c.reduceList[idx].DataPath = append(c.reduceList[idx].DataPath, filepath)
			}
		}
		c.restM--
	} else {
		c.restR--
	}
	c.doneList[uid] = true

	return nil
}

// create a Coordinator.
// main/mrcoordinator.go calls this function.
// nReduce is the number of reduce tasks to use.
// files are the M input files
func MakeCoordinator(files []string, nReduce int) *Coordinator {
	c := Coordinator{}

	// Your code here.

	c.init(len(files), nReduce, files)

	c.server()
	return &c
}
